{
 "cells": [
  {
   "metadata": {},
   "cell_type": "markdown",
   "source": [
    "## U-Net\n",
    "U-Net也是编码器-解码器（Encoder-Decoder）结构，成U型，应用在图像分割领域。\n",
    "\n",
    "U-Net论文：https://arxiv.org/abs/1505.04597"
   ],
   "id": "66cf0c4bf300e010"
  },
  {
   "metadata": {
    "ExecuteTime": {
     "end_time": "2025-08-25T06:02:24.095778Z",
     "start_time": "2025-08-25T06:02:24.088926Z"
    }
   },
   "cell_type": "code",
   "source": [
    "import torch\n",
    "import torch.nn as nn\n",
    "\n",
    "# padding=1使输入和输出尺寸一致，会在本来的图片四周填充0\n",
    "# 本来是26*26的图片，padding后就变成了28*28\n",
    "test_conv = nn.Conv2d(1, 16, 3, padding=1)\n",
    "# 模拟一张单通道的图片\n",
    "test_input = torch.randn(1, 1, 28, 28)\n",
    "\n",
    "print(test_conv(test_input).shape)"
   ],
   "id": "cce6561dea25528d",
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "torch.Size([1, 16, 28, 28])\n"
     ]
    }
   ],
   "execution_count": 32
  },
  {
   "metadata": {
    "ExecuteTime": {
     "end_time": "2025-08-25T05:50:35.563246Z",
     "start_time": "2025-08-25T05:50:35.559667Z"
    }
   },
   "cell_type": "code",
   "source": [
    "import torch\n",
    "import torch.nn as nn\n",
    "\n",
    "test_conv1 = nn.Conv2d(1, 16, 1)\n",
    "test_input1 = torch.randn(1, 1, 28, 28)\n",
    "print(test_conv1(test_input1).shape)"
   ],
   "id": "1f00e51b30e1098a",
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "torch.Size([1, 16, 28, 28])\n"
     ]
    }
   ],
   "execution_count": 21
  },
  {
   "metadata": {
    "ExecuteTime": {
     "end_time": "2025-08-25T05:50:35.589696Z",
     "start_time": "2025-08-25T05:50:35.586657Z"
    }
   },
   "cell_type": "code",
   "source": [
    "test_upsample = nn.Upsample(scale_factor=2, mode='bilinear', align_corners=True)\n",
    "test_input2 = torch.randn(1, 1, 28, 28)\n",
    "print(test_upsample(test_input2).shape)"
   ],
   "id": "7c7c2ebe27452af8",
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "torch.Size([1, 1, 56, 56])\n"
     ]
    }
   ],
   "execution_count": 22
  },
  {
   "metadata": {
    "ExecuteTime": {
     "end_time": "2025-08-25T06:09:08.623134Z",
     "start_time": "2025-08-25T06:09:08.602822Z"
    }
   },
   "cell_type": "code",
   "source": [
    "import torch\n",
    "import torch.nn as nn\n",
    "\n",
    "# 两次卷积操作，卷积核为3，填充为1，使得输入和输出的尺寸一致\n",
    "class DoubleConv(nn.Module):\n",
    "    def __init__(self, in_ch, out_ch):\n",
    "        super(DoubleConv, self).__init__()\n",
    "        self.conv = nn.Sequential(\n",
    "            nn.Conv2d(in_ch, out_ch, 3, padding=1),\n",
    "            nn.BatchNorm2d(out_ch),\n",
    "            nn.ReLU(inplace=True),\n",
    "            nn.Conv2d(out_ch, out_ch, 3, padding=1),\n",
    "            nn.BatchNorm2d(out_ch),\n",
    "            nn.ReLU(inplace=True)\n",
    "        )\n",
    "\n",
    "    def forward(self, x):\n",
    "        return self.conv(x)\n",
    "\n",
    "\n",
    "# 下采样，图片缩小一半\n",
    "class Down(nn.Module):\n",
    "    def __init__(self, in_ch, out_ch):\n",
    "        super(Down, self).__init__()\n",
    "        self.mpconv = nn.Sequential(\n",
    "            nn.MaxPool2d(2),\n",
    "            DoubleConv(in_ch, out_ch)\n",
    "        )\n",
    "\n",
    "    def forward(self, x):\n",
    "        return self.mpconv(x)\n",
    "\n",
    "\n",
    "# 上采样，图片放大两倍\n",
    "class Up(nn.Module):\n",
    "    def __init__(self, in_ch, out_ch):\n",
    "        super(Up, self).__init__()\n",
    "        self.up = nn.Upsample(scale_factor=2, mode='bilinear', align_corners=True)\n",
    "        self.conv = DoubleConv(in_ch, out_ch)\n",
    "\n",
    "    def forward(self, x1, x2):\n",
    "        x1 = self.up(x1)\n",
    "        # 裁剪并拼接跳跃连接\n",
    "        diffY = x2.size()[2] - x1.size()[2]\n",
    "        diffX = x2.size()[3] - x1.size()[3]\n",
    "        pad_left = diffX // 2\n",
    "        pad_right = diffX - pad_left\n",
    "        pad_top = diffY // 2\n",
    "        pad_bottom = diffY - pad_top\n",
    "        x1 = nn.functional.pad(x1, [pad_left, pad_right, pad_top, pad_bottom])\n",
    "        x = torch.cat([x2, x1], dim=1)\n",
    "\n",
    "        return self.conv(x)\n",
    "\n",
    "\n",
    "class OutConv(nn.Module):\n",
    "    def __init__(self, in_ch, out_ch):\n",
    "        super(OutConv, self).__init__()\n",
    "        # 卷积核大小为1，使得输出和输入尺寸一致，但还是变换了图片通道数\n",
    "        # out_ch=n_classes，n_classes=1，所以就是把之前的图片通道数变为1\n",
    "        # 最终整个UNet输入一张图片，输出也是一张图片\n",
    "        self.conv = nn.Conv2d(in_ch, out_ch, 1)\n",
    "\n",
    "    def forward(self, x):\n",
    "        # sigmoid是把像素值变为0-1之间\n",
    "        return torch.sigmoid(self.conv(x))\n",
    "\n",
    "\n",
    "class UNet(nn.Module):\n",
    "    def __init__(self, n_channels=1, n_classes=1):\n",
    "        super(UNet, self).__init__()\n",
    "        self.inc = DoubleConv(n_channels, 64)\n",
    "        self.down1 = Down(64, 128)\n",
    "        self.down2 = Down(128, 256)\n",
    "        self.down3 = Down(256, 512)\n",
    "        self.down4 = Down(512, 512)\n",
    "        self.up1 = Up(512 + 512, 256)  # 通道数：512 (来自下采样) + 512 (上采样输入)，跳跃连接\n",
    "        self.up2 = Up(256 + 256, 128)\n",
    "        self.up3 = Up(128 + 128, 64)\n",
    "        self.up4 = Up(64 + 64, 64)\n",
    "        self.outc = OutConv(64, n_classes)\n",
    "\n",
    "    def forward(self, x):      # 1,28,28\n",
    "        x1 = self.inc(x)       # 64,28,28\n",
    "        x2 = self.down1(x1)    # 128,14,14\n",
    "        x3 = self.down2(x2)    # 256,7,7\n",
    "        x4 = self.down3(x3)    # 512,3,3\n",
    "        x5 = self.down4(x4)    # 512,1,1\n",
    "        x = self.up1(x5, x4)   # 256,3,3\n",
    "        x = self.up2(x, x3)    # 128,7,7\n",
    "        x = self.up3(x, x2)    # 64,14,14\n",
    "        x = self.up4(x, x1)    # 64,28,28\n",
    "        x = self.outc(x)       # 1,28,28\n",
    "        return x"
   ],
   "id": "bc684a7f3b8d5dab",
   "outputs": [],
   "execution_count": 33
  },
  {
   "metadata": {
    "ExecuteTime": {
     "end_time": "2025-08-25T06:09:56.798525Z",
     "start_time": "2025-08-25T06:09:56.762026Z"
    }
   },
   "cell_type": "code",
   "source": [
    "import torch.optim as optim\n",
    "from torch.utils.data import DataLoader\n",
    "from torchvision import datasets, transforms\n",
    "\n",
    "train_dataset = datasets.MNIST(root='./data', train=True, download=True, transform=transforms.ToTensor())\n",
    "test_dataset = datasets.MNIST(root='./data', train=False, download=True, transform=transforms.ToTensor())\n",
    "\n",
    "def collate_fn(batch):\n",
    "    images = torch.stack([item[0] for item in batch])\n",
    "    labels = (images > 0).float()  # 数字区域为1，背景为0，应为mnist中的背后为黑色，像素值等于0\n",
    "    return images, labels\n",
    "\n",
    "train_loader = DataLoader(train_dataset, batch_size=32, shuffle=True, collate_fn=collate_fn)\n",
    "test_loader = DataLoader(test_dataset, batch_size=32, shuffle=False, collate_fn=collate_fn)"
   ],
   "id": "35a391545d4ed9b3",
   "outputs": [],
   "execution_count": 34
  },
  {
   "metadata": {
    "ExecuteTime": {
     "end_time": "2025-08-25T06:14:10.554510Z",
     "start_time": "2025-08-25T06:10:02.468500Z"
    }
   },
   "cell_type": "code",
   "source": [
    "device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')\n",
    "model = UNet(n_channels=1, n_classes=1).to(device)\n",
    "criterion = nn.BCELoss()  # 二元交叉熵损失\n",
    "optimizer = optim.Adam(model.parameters(), lr=1e-3)\n",
    "\n",
    "num_epochs = 10\n",
    "for epoch in range(num_epochs):\n",
    "    model.train()\n",
    "    running_loss = 0.0\n",
    "    for i, (inputs, labels) in enumerate(train_loader):\n",
    "        inputs, labels = inputs.to(device), labels.to(device)\n",
    "\n",
    "        optimizer.zero_grad()\n",
    "        outputs = model(inputs)\n",
    "        loss = criterion(outputs, labels)\n",
    "        loss.backward()\n",
    "        optimizer.step()\n",
    "\n",
    "        running_loss += loss.item()\n",
    "\n",
    "        if i % 100 == 0:\n",
    "            print(f\"Epoch [{epoch + 1}/{num_epochs}], Step [{i}/{len(train_loader)}], Loss: {loss.item():.4f}\")\n",
    "\n",
    "    avg_loss = running_loss / len(train_loader)\n",
    "    print(f\"Epoch [{epoch + 1}/{num_epochs}], Average Loss: {avg_loss:.4f}\")"
   ],
   "id": "9115eb272f22729b",
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch [1/10], Step [0/1875], Loss: 0.6476\n"
     ]
    },
    {
     "ename": "KeyboardInterrupt",
     "evalue": "",
     "output_type": "error",
     "traceback": [
      "\u001B[0;31m---------------------------------------------------------------------------\u001B[0m",
      "\u001B[0;31mKeyboardInterrupt\u001B[0m                         Traceback (most recent call last)",
      "Cell \u001B[0;32mIn[35], line 15\u001B[0m\n\u001B[1;32m     13\u001B[0m optimizer\u001B[38;5;241m.\u001B[39mzero_grad()\n\u001B[1;32m     14\u001B[0m outputs \u001B[38;5;241m=\u001B[39m model(inputs)\n\u001B[0;32m---> 15\u001B[0m loss \u001B[38;5;241m=\u001B[39m \u001B[43mcriterion\u001B[49m(outputs, labels)\n\u001B[1;32m     16\u001B[0m loss\u001B[38;5;241m.\u001B[39mbackward()\n\u001B[1;32m     17\u001B[0m optimizer\u001B[38;5;241m.\u001B[39mstep()\n",
      "Cell \u001B[0;32mIn[35], line 15\u001B[0m\n\u001B[1;32m     13\u001B[0m optimizer\u001B[38;5;241m.\u001B[39mzero_grad()\n\u001B[1;32m     14\u001B[0m outputs \u001B[38;5;241m=\u001B[39m model(inputs)\n\u001B[0;32m---> 15\u001B[0m loss \u001B[38;5;241m=\u001B[39m \u001B[43mcriterion\u001B[49m(outputs, labels)\n\u001B[1;32m     16\u001B[0m loss\u001B[38;5;241m.\u001B[39mbackward()\n\u001B[1;32m     17\u001B[0m optimizer\u001B[38;5;241m.\u001B[39mstep()\n",
      "File \u001B[0;32m_pydevd_bundle/pydevd_cython_darwin_310_64.pyx:1187\u001B[0m, in \u001B[0;36m_pydevd_bundle.pydevd_cython_darwin_310_64.SafeCallWrapper.__call__\u001B[0;34m()\u001B[0m\n",
      "File \u001B[0;32m_pydevd_bundle/pydevd_cython_darwin_310_64.pyx:627\u001B[0m, in \u001B[0;36m_pydevd_bundle.pydevd_cython_darwin_310_64.PyDBFrame.trace_dispatch\u001B[0;34m()\u001B[0m\n",
      "File \u001B[0;32m_pydevd_bundle/pydevd_cython_darwin_310_64.pyx:937\u001B[0m, in \u001B[0;36m_pydevd_bundle.pydevd_cython_darwin_310_64.PyDBFrame.trace_dispatch\u001B[0;34m()\u001B[0m\n",
      "File \u001B[0;32m_pydevd_bundle/pydevd_cython_darwin_310_64.pyx:928\u001B[0m, in \u001B[0;36m_pydevd_bundle.pydevd_cython_darwin_310_64.PyDBFrame.trace_dispatch\u001B[0;34m()\u001B[0m\n",
      "File \u001B[0;32m_pydevd_bundle/pydevd_cython_darwin_310_64.pyx:585\u001B[0m, in \u001B[0;36m_pydevd_bundle.pydevd_cython_darwin_310_64.PyDBFrame.do_wait_suspend\u001B[0;34m()\u001B[0m\n",
      "File \u001B[0;32m/Applications/PyCharm.app/Contents/plugins/python-ce/helpers/pydev/pydevd.py:1220\u001B[0m, in \u001B[0;36mPyDB.do_wait_suspend\u001B[0;34m(self, thread, frame, event, arg, send_suspend_message, is_unhandled_exception)\u001B[0m\n\u001B[1;32m   1217\u001B[0m         from_this_thread\u001B[38;5;241m.\u001B[39mappend(frame_id)\n\u001B[1;32m   1219\u001B[0m \u001B[38;5;28;01mwith\u001B[39;00m \u001B[38;5;28mself\u001B[39m\u001B[38;5;241m.\u001B[39m_threads_suspended_single_notification\u001B[38;5;241m.\u001B[39mnotify_thread_suspended(thread_id, stop_reason):\n\u001B[0;32m-> 1220\u001B[0m     \u001B[38;5;28;43mself\u001B[39;49m\u001B[38;5;241;43m.\u001B[39;49m\u001B[43m_do_wait_suspend\u001B[49m\u001B[43m(\u001B[49m\u001B[43mthread\u001B[49m\u001B[43m,\u001B[49m\u001B[43m \u001B[49m\u001B[43mframe\u001B[49m\u001B[43m,\u001B[49m\u001B[43m \u001B[49m\u001B[43mevent\u001B[49m\u001B[43m,\u001B[49m\u001B[43m \u001B[49m\u001B[43marg\u001B[49m\u001B[43m,\u001B[49m\u001B[43m \u001B[49m\u001B[43msuspend_type\u001B[49m\u001B[43m,\u001B[49m\u001B[43m \u001B[49m\u001B[43mfrom_this_thread\u001B[49m\u001B[43m)\u001B[49m\n",
      "File \u001B[0;32m/Applications/PyCharm.app/Contents/plugins/python-ce/helpers/pydev/pydevd.py:1235\u001B[0m, in \u001B[0;36mPyDB._do_wait_suspend\u001B[0;34m(self, thread, frame, event, arg, suspend_type, from_this_thread)\u001B[0m\n\u001B[1;32m   1232\u001B[0m             \u001B[38;5;28mself\u001B[39m\u001B[38;5;241m.\u001B[39m_call_mpl_hook()\n\u001B[1;32m   1234\u001B[0m         \u001B[38;5;28mself\u001B[39m\u001B[38;5;241m.\u001B[39mprocess_internal_commands()\n\u001B[0;32m-> 1235\u001B[0m         \u001B[43mtime\u001B[49m\u001B[38;5;241;43m.\u001B[39;49m\u001B[43msleep\u001B[49m\u001B[43m(\u001B[49m\u001B[38;5;241;43m0.01\u001B[39;49m\u001B[43m)\u001B[49m\n\u001B[1;32m   1237\u001B[0m \u001B[38;5;28mself\u001B[39m\u001B[38;5;241m.\u001B[39mcancel_async_evaluation(get_current_thread_id(thread), \u001B[38;5;28mstr\u001B[39m(\u001B[38;5;28mid\u001B[39m(frame)))\n\u001B[1;32m   1239\u001B[0m \u001B[38;5;66;03m# process any stepping instructions\u001B[39;00m\n",
      "\u001B[0;31mKeyboardInterrupt\u001B[0m: "
     ]
    }
   ],
   "execution_count": 35
  },
  {
   "metadata": {},
   "cell_type": "code",
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch [1/10], Step [0/1875], Loss: 0.6403\n"
     ]
    },
    {
     "ename": "KeyboardInterrupt",
     "evalue": "",
     "output_type": "error",
     "traceback": [
      "\u001B[0;31m---------------------------------------------------------------------------\u001B[0m",
      "\u001B[0;31mKeyboardInterrupt\u001B[0m                         Traceback (most recent call last)",
      "Cell \u001B[0;32mIn[26], line 31\u001B[0m\n\u001B[1;32m     29\u001B[0m outputs \u001B[38;5;241m=\u001B[39m model(inputs)\n\u001B[1;32m     30\u001B[0m loss \u001B[38;5;241m=\u001B[39m criterion(outputs, labels)\n\u001B[0;32m---> 31\u001B[0m \u001B[43mloss\u001B[49m\u001B[38;5;241;43m.\u001B[39;49m\u001B[43mbackward\u001B[49m\u001B[43m(\u001B[49m\u001B[43m)\u001B[49m\n\u001B[1;32m     32\u001B[0m optimizer\u001B[38;5;241m.\u001B[39mstep()\n\u001B[1;32m     34\u001B[0m running_loss \u001B[38;5;241m+\u001B[39m\u001B[38;5;241m=\u001B[39m loss\u001B[38;5;241m.\u001B[39mitem()\n",
      "File \u001B[0;32m~/miniconda3/envs/mini-gpt/lib/python3.10/site-packages/torch/_tensor.py:648\u001B[0m, in \u001B[0;36mTensor.backward\u001B[0;34m(self, gradient, retain_graph, create_graph, inputs)\u001B[0m\n\u001B[1;32m    638\u001B[0m \u001B[38;5;28;01mif\u001B[39;00m has_torch_function_unary(\u001B[38;5;28mself\u001B[39m):\n\u001B[1;32m    639\u001B[0m     \u001B[38;5;28;01mreturn\u001B[39;00m handle_torch_function(\n\u001B[1;32m    640\u001B[0m         Tensor\u001B[38;5;241m.\u001B[39mbackward,\n\u001B[1;32m    641\u001B[0m         (\u001B[38;5;28mself\u001B[39m,),\n\u001B[0;32m   (...)\u001B[0m\n\u001B[1;32m    646\u001B[0m         inputs\u001B[38;5;241m=\u001B[39minputs,\n\u001B[1;32m    647\u001B[0m     )\n\u001B[0;32m--> 648\u001B[0m \u001B[43mtorch\u001B[49m\u001B[38;5;241;43m.\u001B[39;49m\u001B[43mautograd\u001B[49m\u001B[38;5;241;43m.\u001B[39;49m\u001B[43mbackward\u001B[49m\u001B[43m(\u001B[49m\n\u001B[1;32m    649\u001B[0m \u001B[43m    \u001B[49m\u001B[38;5;28;43mself\u001B[39;49m\u001B[43m,\u001B[49m\u001B[43m \u001B[49m\u001B[43mgradient\u001B[49m\u001B[43m,\u001B[49m\u001B[43m \u001B[49m\u001B[43mretain_graph\u001B[49m\u001B[43m,\u001B[49m\u001B[43m \u001B[49m\u001B[43mcreate_graph\u001B[49m\u001B[43m,\u001B[49m\u001B[43m \u001B[49m\u001B[43minputs\u001B[49m\u001B[38;5;241;43m=\u001B[39;49m\u001B[43minputs\u001B[49m\n\u001B[1;32m    650\u001B[0m \u001B[43m\u001B[49m\u001B[43m)\u001B[49m\n",
      "File \u001B[0;32m~/miniconda3/envs/mini-gpt/lib/python3.10/site-packages/torch/autograd/__init__.py:353\u001B[0m, in \u001B[0;36mbackward\u001B[0;34m(tensors, grad_tensors, retain_graph, create_graph, grad_variables, inputs)\u001B[0m\n\u001B[1;32m    348\u001B[0m     retain_graph \u001B[38;5;241m=\u001B[39m create_graph\n\u001B[1;32m    350\u001B[0m \u001B[38;5;66;03m# The reason we repeat the same comment below is that\u001B[39;00m\n\u001B[1;32m    351\u001B[0m \u001B[38;5;66;03m# some Python versions print out the first line of a multi-line function\u001B[39;00m\n\u001B[1;32m    352\u001B[0m \u001B[38;5;66;03m# calls in the traceback and some print out the last line\u001B[39;00m\n\u001B[0;32m--> 353\u001B[0m \u001B[43m_engine_run_backward\u001B[49m\u001B[43m(\u001B[49m\n\u001B[1;32m    354\u001B[0m \u001B[43m    \u001B[49m\u001B[43mtensors\u001B[49m\u001B[43m,\u001B[49m\n\u001B[1;32m    355\u001B[0m \u001B[43m    \u001B[49m\u001B[43mgrad_tensors_\u001B[49m\u001B[43m,\u001B[49m\n\u001B[1;32m    356\u001B[0m \u001B[43m    \u001B[49m\u001B[43mretain_graph\u001B[49m\u001B[43m,\u001B[49m\n\u001B[1;32m    357\u001B[0m \u001B[43m    \u001B[49m\u001B[43mcreate_graph\u001B[49m\u001B[43m,\u001B[49m\n\u001B[1;32m    358\u001B[0m \u001B[43m    \u001B[49m\u001B[43minputs\u001B[49m\u001B[43m,\u001B[49m\n\u001B[1;32m    359\u001B[0m \u001B[43m    \u001B[49m\u001B[43mallow_unreachable\u001B[49m\u001B[38;5;241;43m=\u001B[39;49m\u001B[38;5;28;43;01mTrue\u001B[39;49;00m\u001B[43m,\u001B[49m\n\u001B[1;32m    360\u001B[0m \u001B[43m    \u001B[49m\u001B[43maccumulate_grad\u001B[49m\u001B[38;5;241;43m=\u001B[39;49m\u001B[38;5;28;43;01mTrue\u001B[39;49;00m\u001B[43m,\u001B[49m\n\u001B[1;32m    361\u001B[0m \u001B[43m\u001B[49m\u001B[43m)\u001B[49m\n",
      "File \u001B[0;32m~/miniconda3/envs/mini-gpt/lib/python3.10/site-packages/torch/autograd/graph.py:824\u001B[0m, in \u001B[0;36m_engine_run_backward\u001B[0;34m(t_outputs, *args, **kwargs)\u001B[0m\n\u001B[1;32m    822\u001B[0m     unregister_hooks \u001B[38;5;241m=\u001B[39m _register_logging_hooks_on_whole_graph(t_outputs)\n\u001B[1;32m    823\u001B[0m \u001B[38;5;28;01mtry\u001B[39;00m:\n\u001B[0;32m--> 824\u001B[0m     \u001B[38;5;28;01mreturn\u001B[39;00m \u001B[43mVariable\u001B[49m\u001B[38;5;241;43m.\u001B[39;49m\u001B[43m_execution_engine\u001B[49m\u001B[38;5;241;43m.\u001B[39;49m\u001B[43mrun_backward\u001B[49m\u001B[43m(\u001B[49m\u001B[43m  \u001B[49m\u001B[38;5;66;43;03m# Calls into the C++ engine to run the backward pass\u001B[39;49;00m\n\u001B[1;32m    825\u001B[0m \u001B[43m        \u001B[49m\u001B[43mt_outputs\u001B[49m\u001B[43m,\u001B[49m\u001B[43m \u001B[49m\u001B[38;5;241;43m*\u001B[39;49m\u001B[43margs\u001B[49m\u001B[43m,\u001B[49m\u001B[43m \u001B[49m\u001B[38;5;241;43m*\u001B[39;49m\u001B[38;5;241;43m*\u001B[39;49m\u001B[43mkwargs\u001B[49m\n\u001B[1;32m    826\u001B[0m \u001B[43m    \u001B[49m\u001B[43m)\u001B[49m  \u001B[38;5;66;03m# Calls into the C++ engine to run the backward pass\u001B[39;00m\n\u001B[1;32m    827\u001B[0m \u001B[38;5;28;01mfinally\u001B[39;00m:\n\u001B[1;32m    828\u001B[0m     \u001B[38;5;28;01mif\u001B[39;00m attach_logging_hooks:\n",
      "\u001B[0;31mKeyboardInterrupt\u001B[0m: "
     ]
    }
   ],
   "execution_count": 26,
   "source": [
    "import matplotlib.pyplot as plt\n",
    "\n",
    "model.eval()\n",
    "with torch.no_grad():\n",
    "    data_iter = iter(test_loader)\n",
    "    images, labels = next(data_iter)\n",
    "    images, labels = images[:4].to(device), labels[:4]\n",
    "\n",
    "    outputs = model(images)\n",
    "    preds = (outputs > 0.5).float()  # 阈值化为二值图\n",
    "\n",
    "    # 转为 NumPy 以便可视化\n",
    "    images = images.cpu().numpy()\n",
    "    labels = labels.cpu().numpy()\n",
    "    preds = preds.cpu().numpy()\n",
    "\n",
    "    # 绘图\n",
    "    fig, axes = plt.subplots(3, 4, figsize=(10, 6))\n",
    "    for i in range(4):\n",
    "        axes[0, i].imshow(images[i, 0], cmap='gray')\n",
    "        axes[0, i].set_title(\"Input Image\")\n",
    "        axes[0, i].axis('off')\n",
    "\n",
    "        axes[1, i].imshow(labels[i, 0], cmap='gray')\n",
    "        axes[1, i].set_title(\"True Mask\")\n",
    "        axes[1, i].axis('off')\n",
    "\n",
    "        axes[2, i].imshow(preds[i, 0], cmap='gray')\n",
    "        axes[2, i].set_title(\"Predicted Mask\")\n",
    "        axes[2, i].axis('off')\n",
    "\n",
    "    plt.tight_layout()\n",
    "    plt.show()"
   ],
   "id": "d3e106890fd5c6bf"
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 2
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython2",
   "version": "2.7.6"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
